import numpy as np
import torch
import torch.nn as nn
from scipy.stats import spearmanr
from src.diff_models_table import diff_CSDI
import yaml
import torch.nn.functional as F
from torch.nn.functional import huber_loss  # 添加这行
class CSDI_base(nn.Module):
    def __init__(self, target_dim, config, device, max_values=None, min_values=None):
        super().__init__()
        #设备与维度设置
        self.device = device #计算设备（cpu/GPU）
        self.target_dim = target_dim #目标特征维度（比如传感器数量）
        print("target_dim",self.target_dim)

        #嵌入维度配置
        # 从配置获取嵌入维度参数
        self.emb_time_dim = config["model"]["timeemb"] #时间编码维度
        self.emb_feature_dim = config["model"]["featureemb"] #特征编码维度

        # 模型模式设置
        self.is_unconditional = config["model"]["is_unconditional"] #是否是无条件生成
        self.target_strategy = config["model"]["target_strategy"] #目标策略（未实际使用）

        # 计算总嵌入维度 = 时间嵌入维度+特征嵌入维度
        self.emb_total_dim = self.emb_time_dim + self.emb_feature_dim
        if self.is_unconditional == False:
            self.emb_total_dim += 1  # for conditional mask 条件模式下增加mask通道1维用于传递mask信息

        # 特征嵌入层（每个特征ID映射为向量）
        self.embed_layer = nn.Embedding(
            num_embeddings=self.target_dim, # 嵌入字典大小=特征数
            embedding_dim=self.emb_feature_dim  # 每个特征的嵌入维度
        )

        # 扩散模型配置
        config_diff = config["diffusion"] #调用扩散模型的配置
        config_diff["side_dim"] = self.emb_total_dim # 设置侧信息维度 将总嵌入维度加入配置
        input_dim = 1 if self.is_unconditional == True else 2 # 输入通道数 无条件输入1通道，条件输入2通道
        self.diffmodel = diff_CSDI(config_diff, input_dim) # 初始化核心扩散模型

        #噪声调度设置
        # parameters for diffusion models
        self.num_steps = config_diff["num_steps"] # 扩散总步数（默认50） 150
        if config_diff["schedule"] == "quad": # 二次方调度
            self.beta = (
                np.linspace(
                    config_diff["beta_start"] ** 0.5, # β起始值（默认0.0001）
                    config_diff["beta_end"] ** 0.5, # β结束值（默认0.5）
                    self.num_steps,
                )
                ** 2
            )# 先线性插值再平方
        elif config_diff["schedule"] == "linear": # 线性调度
            self.beta = np.linspace(
                config_diff["beta_start"], config_diff["beta_end"], self.num_steps
            ) #直接线性插值

        #α相关计算
        self.alpha_hat = 1 - self.beta #β是噪声添加比例，那么α_hat就是信号保留比例
        self.alpha = np.cumprod(self.alpha_hat) # α是α_hat的累积乘积
        self.alpha_torch = (
            torch.tensor(self.alpha).float().to(self.device).unsqueeze(1).unsqueeze(1)
        )# 转换为(B,1,1)形状的张量便于广播计算

        # # 新增最大值和最小值约束
        # self.max_values = torch.tensor(max_values * 1.1, device=self.device)  # 放宽 10%  这里可以做一个更改看精度变化
        # self.min_values = torch.tensor(min_values * 0.9, device=self.device)  # 放宽 10%
        #
        # # 如果 max_values 和 min_values 是 numpy 数组，将其转换为 tensor 冗余检查可删除
        # if self.max_values is not None:
        #     self.max_values = torch.tensor(self.max_values, device=self.device)
        # if self.min_values is not None:
        #     self.min_values = torch.tensor(self.min_values, device=self.device)

        # +++ 修改这里：更智能的边界处理 +++
        if max_values is not None and min_values is not None:
            # 如果提供了max/min值，直接使用（放宽10%）
            self.max_values = torch.tensor(max_values * 1.1, device=self.device)
            self.min_values = torch.tensor(min_values * 0.9, device=self.device)
            print(f"使用提供的边界值: min={self.min_values}, max={self.max_values}")
        else:
            # 如果没提供，初始化为None，等待从数据中学习
            self.max_values = None
            self.min_values = None
            self.learned_max = None
            self.learned_min = None
            print("未提供边界值，将从训练数据中学习")
# 更新边界最值
    def update_data_bounds(self, observed_data, observed_mask):
        """
        从观测数据中学习每个特征的最大最小值
        observed_data: (B, K, L) 观测数据
        observed_mask: (B, K, L) 观测掩码
        """
        if self.max_values is not None:
            return  # 如果已经提供了边界值，就不需要学习

        with torch.no_grad():
            # 只考虑观测到的数据点
            valid_data = observed_data * observed_mask
            valid_mask = observed_mask > 0

            # 计算每个特征的最大最小值
            current_max = []
            current_min = []

            for k in range(observed_data.shape[1]):  # 遍历每个特征
                feature_data = valid_data[:, k, :]
                feature_mask = valid_mask[:, k, :]

                if feature_mask.any():
                    # 只考虑有观测值的点
                    masked_data = feature_data[feature_mask]
                    current_max.append(masked_data.max().item())
                    current_min.append(masked_data.min().item())
                else:
                    # 如果没有观测值，使用默认范围
                    current_max.append(1.0)
                    current_min.append(0.0)

            # 转换为tensor
            current_max = torch.tensor(current_max, device=self.device)
            current_min = torch.tensor(current_min, device=self.device)

            # 更新学习到的边界（逐步扩大范围）
            if self.learned_max is None:
                self.learned_max = current_max * 1.1  # 初始放宽10%
                self.learned_min = current_min * 0.9
            else:
                # 逐步调整边界，包含新看到的数据
                self.learned_max = torch.maximum(self.learned_max, current_max * 1.1)
                self.learned_min = torch.minimum(self.learned_min, current_min * 0.9)

            print(f"学习到的数据边界 - 最小值: {self.learned_min}")
            print(f"学习到的数据边界 - 最大值: {self.learned_max}")
    #时间编码方法
    def time_embedding(self, pos, d_model=128):
        pe = torch.zeros(pos.shape[0], pos.shape[1], d_model).to(self.device)
        # 初始化位置编码矩阵 (B,L,d_model)
        position = pos.unsqueeze(2)# (B,L,1)
        div_term = 1 / torch.pow(
            10000.0, torch.arange(0, d_model, 2).to(self.device) / d_model
        )# 计算除数项 [10000^(2i/d_model)]
        pe[:, :, 0::2] = torch.sin(position * div_term)# 偶数位sin
        pe[:, :, 1::2] = torch.cos(position * div_term)# 奇数位cos
        return pe# 输出形状 (B,L,d_model)

    #随机掩码生成这段代码实现了一种随机掩码生成（Random Masking）方法，主要用于处理部分观测数据（如时间序列中的缺失值）。
    # 核心功能是：在已知的观测点（observed_mask标记为1的位置）中，随机选择一部分点作为“保留点”，其余观测点则被掩码（视为缺失）
    def get_randmask(self, observed_mask):
        # 在 observed_mask=1 的位置生成均匀分布随机数（[0, 1)），其余位置为 0
        rand_for_mask = torch.rand_like(observed_mask) * observed_mask
        rand_for_mask = rand_for_mask.reshape(len(rand_for_mask), -1)# 展平 (B, K*L)

        for i in range(len(observed_mask)):
            sample_ratio = 0.3  # np.random.rand()  # missing ratio
            num_observed = observed_mask[i].sum().item()#对每个样本，根据 observed_mask 统计有效观测点数 num_observed
            num_masked = round(num_observed * sample_ratio)#计算需要掩码的点数 num_masked
            #使用 topk 找到随机数中最大的 num_masked 个位置（即随机选择），将其值设为 -1
            rand_for_mask[i][rand_for_mask[i].topk(num_masked).indices] = -1
            # 真正的掩码，真实缺失的以及设置观测到80%缺失的为false,20%观测到的为true
        cond_mask = (rand_for_mask > 0).reshape(observed_mask.shape).float()
        return cond_mask #返回掩码 浮点型

    #侧信息生成 侧信息通常包括时间嵌入、特征嵌入（Feature Embedding）和条件掩码（Conditional Mask），
    # 它们为模型提供额外的上下文信息（如时间戳、变量类型、缺失模式等
    def get_side_info(self, observed_tp, cond_mask):
        #输入：observed_tp观测时间点（B，L）  cond_mask条件掩码(B,K,L) 输出side_info：侧信息张量，形状为 (B, D, K, L)D是嵌入总维度
        B, K, L = cond_mask.shape

        # 时间戳 observed_tp 转换为 (B, L, emb_time_dim) 的向量(B,L,emb)
        time_embed = self.time_embedding(observed_tp, self.emb_time_dim)
        #通过unsqueeze(2)和expand将时间嵌入扩展到每个特征维度（K），(B, L, K, emb_time_dim)。
        time_embed = time_embed.unsqueeze(2).expand(-1, -1, K, -1)

        #特征嵌入
        feature_embed = self.embed_layer(
            torch.arange(self.target_dim).to(self.device)
        )  # (K,emb)
        feature_embed = feature_embed.unsqueeze(0).unsqueeze(0).expand(B, L, -1, -1)# 扩展维度(B,L,K,emb)



        #拼接时间和特征嵌入
        side_info = torch.cat([time_embed, feature_embed], dim=-1)  # (B,L,K,emb_time+emb_feature)
        side_info = side_info.permute(0, 3, 2, 1)  # (B,D,K,L) D=emb_time+emb_feature 通过 permute 调整维度顺序

        #条件掩码处理
        if self.is_unconditional == False:
            side_mask = cond_mask.unsqueeze(1)  # (B,1,K,L)
            side_info = torch.cat([side_info, side_mask], dim=1) # (B, D+1, K, L)，D+1=总嵌入维度（时间 + 特征 + 掩码）

        return side_info

    #在模型验证（或测试）时，通过多步（num_steps）前向和反向计算来评估模型的平均性能。
    def calc_loss_valid(
        self, observed_data, cond_mask, observed_mask, side_info, is_train
    ):
        loss_sum = 0
        # In validation, perform T steps forward and backward.
        for t in range(self.num_steps):
            loss = self.calc_loss(
                observed_data, cond_mask, observed_mask, side_info, is_train, set_t=t
            )
            loss_sum += loss.detach() #分离损失值，避免梯度计算（验证阶段无需反向传播）
        return loss_sum / self.num_steps

    def masked_pcc(self, x, mask):
        """
        计算掩码后的特征间 Pearson 相关系数矩阵（适配 L=1 的表格数据）
        输入:
            x: (B, K, 1) - 输入数据（L=1）
            mask: (B, K, 1) - 掩码
        输出:
            corr: (B, K, K) - 每个 batch 的特征间 PCC 矩阵
        """
        B, K, L = x.shape
        assert L == 1, "此函数仅支持 L=1 的表格数据"

        # 移除 L 维度并应用掩码
        x_masked = x.squeeze(-1) * mask.squeeze(-1)  # (B, K)

        # 计算均值（跨样本）
        mean = (x_masked.sum(dim=0)) / (mask.squeeze(-1).sum(dim=0) + 1e-6)  # (K,)
        x_centered = (x_masked - mean.unsqueeze(0)) * mask.squeeze(-1)  # (B, K)

        # 计算协方差矩阵 (K, K)
        numerator = torch.matmul(x_centered.T, x_centered)  # (K, K)

        # 计算标准差乘积
        x_std = torch.sqrt((x_centered ** 2).sum(dim=0)) + 1e-6  # (K,)
        denominator = torch.outer(x_std, x_std)  # (K, K)

        corr = numerator / denominator
        return corr.unsqueeze(0)  # 返回 (1, K, K) 以兼容原有接口

    def compute_pcc(self, x, y):
        """
        计算两组特征间的 Pearson 相关系数矩阵（适配 L=1）
        输入:
            x, y: (B, K, 1) - 预测值和目标值
        输出:
            corr: (B, K, K) - 相关系数矩阵
        """
        B, K, L = x.shape
        assert L == 1, "此函数仅支持 L=1 的表格数据"

        # 移除 L 维度
        x_flat = x.squeeze(-1)  # (B, K)
        y_flat = y.squeeze(-1)  # (B, K)

        # 均值归一化
        x_mean = x_flat.mean(dim=0, keepdim=True)  # (1, K)
        y_mean = y_flat.mean(dim=0, keepdim=True)  # (1, K)
        x_centered = x_flat - x_mean  # (B, K)
        y_centered = y_flat - y_mean  # (B, K)

        # 协方差矩阵 (K, K)
        numerator = torch.matmul(x_centered.T, y_centered)  # (K, K)

        # 标准差乘积
        x_std = torch.sqrt((x_centered ** 2).sum(dim=0)) + 1e-6  # (K,)
        y_std = torch.sqrt((y_centered ** 2).sum(dim=0)) + 1e-6  # (K,)
        denominator = torch.outer(x_std, y_std)  # (K, K)

        corr = numerator / denominator
        return corr.unsqueeze(0)  # (1, K, K)

    def compute_scc_vector(self, x, ref_target, mask):
        """
        计算每个特征与目标列的跨样本Spearman相关系数（适配L=1的表格数据）
        输入:
            x: (B, K, 1) - 预测值
            ref_target: (B, K, 1) - 真实值（最后一列为目标）
            mask: (B, K, 1) - 掩码
        输出:
            scc_vec: (K,) - 每个特征与目标的SCC（跨样本）
        """
        B, K, L = x.shape
        assert L == 1, "仅支持L=1的表格数据"

        # 提取目标列（形状 (B,)）
        target = ref_target[:, -1, 0].detach().cpu().numpy()  # 假设目标在最后一列

        scc_vec = torch.zeros(K).to(self.device)

        for i in range(K):
            # 收集当前特征的所有有效样本值
            feat_i = x[:, i, 0].detach().cpu().numpy()  # (B,)
            mask_i = mask[:, i, 0].cpu().numpy() > 0  # (B,)

            valid_feat = feat_i[mask_i]
            valid_target = target[mask_i]

            if len(valid_feat) >= 2:  # 至少需要2个样本计算SCC
                scc, _ = spearmanr(valid_feat, valid_target)
                scc_vec[i] = 0.0 if np.isnan(scc) else scc

        return scc_vec.unsqueeze(0)  # 返回 (1, K) 兼容原有接口
    def calc_loss(
            self, observed_data, cond_mask, observed_mask, side_info, is_train, set_t=-1
    ):
        B, K, L = observed_data.shape


        # ========== 时间步调度 ==========
        if is_train != 1:
            t = (torch.ones(B) * set_t).long().to(self.device)
        else:
            t = torch.randint(0, self.num_steps, [B]).to(self.device)

        # ========== 扩散采样 ==========
        current_alpha = self.alpha_torch[t]
        noise = torch.randn_like(observed_data)
        noisy_data = (current_alpha ** 0.5) * observed_data + (1.0 - current_alpha) ** 0.5 * noise
        total_input = self.set_input_to_diffmodel(noisy_data, observed_data, cond_mask)


        predicted = self.diffmodel(total_input, side_info, t)

        # ========== 目标区域验证 ==========
        target_mask = observed_mask - cond_mask
        # DEBUG: 检查目标区域是否有效
        # print(f"目标区域有效点数: {target_mask.sum().item()}/{B * K * L}")  # 应 >0
        if target_mask.sum() == 0:
            raise ValueError("目标区域全被掩码，无法计算PCC/SCC！")

        residual = (noise - predicted) * target_mask
        num_eval = target_mask.sum()

        # ========== 最值约束 ==========
        if self.max_values is not None:
            max_violation = torch.clamp(predicted - self.max_values, min=0.0)
            min_violation = torch.clamp(self.min_values - predicted, min=0.0)
            constraint_loss = huber_loss(max_violation, torch.zeros_like(max_violation)) + \
                              huber_loss(min_violation, torch.zeros_like(min_violation))
        else:
            constraint_loss = torch.tensor(0.0, device=self.device)

        # ========== PCC结构约束 ==========
        if K >= 2:
            true_corr = self.masked_pcc(observed_data, observed_mask)
            pred_corr = self.compute_pcc(predicted, predicted)  # 注意: 这里用predicted自身计算可能有问题

            structure_constraint_loss = F.mse_loss(pred_corr, true_corr)
        else:
            structure_constraint_loss = torch.tensor(0.0, device=self.device)

        # ========== SCC目标约束 ==========
        if K >= 2:
            # DEBUG: 验证SCC计算
            true_scc = self.compute_scc_vector(observed_data, observed_data, observed_mask)
            pred_scc = self.compute_scc_vector(predicted, observed_data, observed_mask)


            scc_loss = F.mse_loss(pred_scc, true_scc)
        else:
            scc_loss = torch.tensor(0.0, device=self.device)

        # ========== 总损失 ==========
        mse_loss = (residual ** 2).sum() / (num_eval if num_eval > 0 else 1)
        # loss = mse_loss + 0.1 * constraint_loss +0.1 * structure_constraint_loss+ 0.1 * scc_loss  # 暂不启用PCC约束
        loss = mse_loss




        return loss


    #输入预处理
    def set_input_to_diffmodel(self, noisy_data, observed_data, cond_mask):
        if self.is_unconditional == True: #无条件生成
            #直接将noisy_data增加一个维度作为输入  这种模式下模型不利用任何条件信息，完全从噪声生成数据
            total_input = noisy_data.unsqueeze(1)  # (B,1,K,L)
        else: #条件生成
            cond_obs = (cond_mask * observed_data).unsqueeze(1) #用cond_mask从observed_data中提取已知部分
            noisy_target = ((1 - cond_mask) * noisy_data).unsqueeze(1) #用(1 - cond_mask)从noisy_data中提取需要预测的部分
            #将两部分在第二维进行拼接
            total_input = torch.cat([cond_obs, noisy_target], dim=1)  # (B,2,K,L)

        return total_input

    #扩散模型的填充过程，从部分观测数据生成完整数据
    def impute(self, observed_data, cond_mask, side_info, n_samples, noisy_cond_history=None):
        #n_samples 为每个样本生成的候选数
        B, K, L = observed_data.shape
        imputed_samples = torch.zeros(B, n_samples, K, L).to(self.device) #初始化输出容器

        for i in range(n_samples): # 生成n_samples个样本
            current_sample = torch.randn_like(observed_data)  # 从高斯噪声开始

            for t in range(self.num_steps - 1, -1, -1):# 逆向扩散过程从最大噪声步(num_steps-1)逆向回到0步逐步去除噪声
                #构造模型输入 无条件生成模式
                if self.is_unconditional == True:
                    diff_input = (
                            cond_mask * noisy_cond_history[t]
                            + (1.0 - cond_mask) * current_sample
                    )#混合噪声条件历史和当前噪声样本  条件区域使用噪声历史，目标区域使用当前噪声样本
                    diff_input = diff_input.unsqueeze(1)  # (B,1,K,L)
                else: #条件生成模型
                    cond_obs = (cond_mask * observed_data).unsqueeze(1) #通道1 原始观测的条件部分（干净数据）
                    noisy_target = ((1 - cond_mask) * current_sample).unsqueeze(1) #通道2 当前噪声样本的目标部分
                    diff_input = torch.cat([cond_obs, noisy_target], dim=1)  # (B,2,K,L)双通道输入
                predicted = self.diffmodel(
                    diff_input, side_info, torch.tensor([t]).to(self.device)
                )  # (B,K,L) 扩散模型使用构造的输入、辅助信息和当前时间步来预测当前步的噪声
                # print(f"t={t} | current_sample: {current_sample.shape} | predicted: {predicted.shape}")
                #更新当前样本  计算逆向过程的系数 使用DDPM采样更新规则
                coeff1 = 1 / self.alpha_hat[t] ** 0.5
                coeff2 = (1 - self.alpha_hat[t]) / (1 - self.alpha[t]) ** 0.5

                current_sample = coeff1 * (current_sample - coeff2 * predicted) # 更新当前样本
                #添加过程噪声
                if t > 0: #如果不是最后一步（t=0） 添加过程噪声
                    noise = torch.randn_like(current_sample) # 生成与当前样本同形的标准高斯噪声
                    sigma = (
                                    (1.0 - self.alpha[t - 1]) / (1.0 - self.alpha[t]) * self.beta[t]
                            ) ** 0.5 #计算噪声强度 确保逆向过程与正向一致
                    current_sample += sigma * noise #将缩放后的噪声添加到当前样本中，使生成过程具有随机性

            # 硬性约束 截断
            if self.max_values is not None:
                # 修改前：current_sample = torch.clamp(current_sample, ...)
                # 修改后：显式保持维度
                current_sample = torch.clamp(current_sample,
                                             min=self.min_values.unsqueeze(-1),  # 确保形状兼容
                                             max=self.max_values.unsqueeze(-1))

            imputed_samples[:, i] = current_sample.detach() #存储结果


        return imputed_samples

    def forward(self, batch, is_train=1):
        # 确保所有输入张量都在正确设备上
        batch = {k: v.to(self.device) if torch.is_tensor(v) else v
                 for k, v in batch.items()}
        (
            observed_data,
            observed_mask,
            observed_tp,
            gt_mask,
            for_pattern_mask,
            _,
        ) = self.process_data(batch)

        # In testing, using `gt_mask` (generated with fixed missing rate).
        if is_train == 0:
            cond_mask = gt_mask #固定保证可复现
        # In training, generate random mask
        else:
            cond_mask = self.get_randmask(observed_mask)# 生成随机掩码

        side_info = self.get_side_info(observed_tp, cond_mask)

        loss_func = self.calc_loss if is_train == 1 else self.calc_loss_valid
        return loss_func(observed_data, cond_mask, observed_mask, side_info, is_train)

    #用于在测试集上评估扩散模型性能的核心方法，它会生成多个可能的填补样本并与真实数据进行比较
    def evaluate(self, batch, n_samples):
        (
            observed_data,
            observed_mask,
            observed_tp,# 时间点坐标 (B, L)
            gt_mask,# 预设的固定掩码 (B, K, L)
            _,
            cut_length,
        ) = self.process_data(batch)

        with torch.no_grad():# 禁用梯度计算
            cond_mask = gt_mask
            target_mask = observed_mask - cond_mask #目标区域
            side_info = self.get_side_info(observed_tp, cond_mask)

            samples = self.impute(observed_data, cond_mask, side_info, n_samples) #生成填补样本


        return samples, observed_data, target_mask, observed_mask, observed_tp

#是 CSDI_base 的子类，专为表格型数据（单变量时间序列）设计
class TabCSDI(CSDI_base):
    def __init__(self, config, device, target_dim, max_values=None, min_values=None):
        super(TabCSDI, self).__init__(target_dim, config, device, max_values=max_values, min_values=min_values)


    def process_data(self, batch):
        # 原始输入数据为 (B, K)


        # 1. 添加伪时间维度 (B, K, 1)
        observed_data = batch["observed_data"].unsqueeze(-1).to(self.device).float()  # (B, K, 1)
        observed_mask = batch["observed_mask"].unsqueeze(-1).to(self.device).float()  # (B, K, 1)
        gt_mask = batch["gt_mask"].unsqueeze(-1).to(self.device).float()  # (B, K, 1)

        # 2. 构造伪时间点 (B, L=1)，时间值可以为 0 或任意固定值
        B, K, L = observed_data.shape
        observed_tp = torch.zeros(B, L).to(self.device).float()  # (B, 1)

        # 3. 兼容父类接口
        cut_length = torch.zeros(B).long().to(self.device)  # (B,)
        for_pattern_mask = observed_mask  # (B, K, L)

        return (
            observed_data,  # (B, K, L)
            observed_mask,  # (B, K, L)
            observed_tp,  # (B, L)
            gt_mask,  # (B, K, L)
            for_pattern_mask,  # (B, K, L)
            cut_length  # (B,)
        )

    # def process_data(self, batch):
    #     print("observes_data.shape",batch["observed_data"].shape)
    #     # Insert K=1 axis. All mask now with shape (B, 1, L). # 1. 观测数据升维：(B,L) -> (B,1,L)
    #     observed_data = batch["observed_data"][:, np.newaxis, :] # 插入K=1维度
    #     observed_data = observed_data.to(self.device).float() # 确保设备和类型都正确
    #
    #     # 2. 观测掩码升维：(B,L) -> (B,1,L)
    #     observed_mask = batch["observed_mask"][:, np.newaxis, :]
    #     observed_mask = observed_mask.to(self.device).float()# 确保在正确设备上
    #     # 3. 时间点数据（不变形状）
    #     observed_tp = batch["timepoints"].to(self.device).float()
    #     # 4. 真实掩码升维（用于评估）
    #     gt_mask = batch["gt_mask"][:, np.newaxis, :]
    #     gt_mask = gt_mask.to(self.device).float()
    #     # 5. 其他参数初始化
    #     cut_length = torch.zeros(len(observed_data)).long().to(self.device) # 未使用（兼容父类）
    #     for_pattern_mask = observed_mask  # 直接复用观测掩码
    #
    #     return (
    #         observed_data,
    #         observed_mask,
    #         observed_tp,
    #         gt_mask,
    #         for_pattern_mask,
    #         cut_length,
    #     )